import fasttext.FastText as fasttext
import os
from utils.word.core_word import WordAnalyzer


class FastTextClassifier:
    '''
    FastText分类器

    Attributes:

    1. model，_FastText，应使用train_model初始化

    2. wa，WordAnalyzer，若用户传入得文本没有分词，则需要调用
    '''

    def __init__(self):
        self.model = None
        self.wa = None

    def train_model(self, ipt=None, opt=None, model='', dim=100, epoch=5, lr=0.1, loss='softmax'):
        '''
        训练FastText模型

        当model指向已有模型时，则直接使用模型
        '''

        if os.path.isfile(model):
            self.model = fasttext.load_model(model)
        else:
            self.model = fasttext.train_supervised(ipt, label='__label__', dim=dim, epoch=epoch,
                                                   lr=lr, wordNgrams=2, loss=loss)
            """
            训练一个监督模型, 返回一个模型对象

            @param input:           训练数据文件路径
            @param lr:              学习率
            @param dim:             向量维度
            @param ws:              cbow模型时使用
            @param epoch:           次数
            @param minCount:        词频阈值, 小于该值在初始化时会过滤掉
            @param minCountLabel:   类别阈值，类别小于该值初始化时会过滤掉
            @param minn:            构造subword时最小char个数
            @param maxn:            构造subword时最大char个数
            @param neg:             负采样
            @param wordNgrams:      n-gram个数
            @param loss:            损失函数类型, softmax, ns: 负采样, hs: 分层softmax
            @param bucket:          词扩充大小, [A, B]: A语料中包含的词向量, B不在语料中的词向量
            @param thread:          线程个数, 每个线程处理输入数据的一段, 0号线程负责loss输出
            @param lrUpdateRate:    学习率更新
            @param t:               负采样阈值
            @param label:           类别前缀
            @param verbose:         ??
            @param pretrainedVectors: 预训练的词向量文件路径, 如果word出现在文件夹中初始化不再随机
            @return model object
            """
            self.model.save_model(opt)

    def quantize_model(self, input=None, qout=False, cutoff=0, epoch=None, lr=None, retrain=False):
        """
        将训练模型压缩

        如果retrain=False，则按照默认压缩

        否则，需要输入训练集和一系列训练参数
        """

        if self.model is None:
            return

        if retrain == False:
            self.model.quantize()
        else:
            self.model.quantize(input=input, qout=qout,
                                cutoff=cutoff, retrain=True, epoch=epoch, lr=lr)
        self.model.save_model("quantized_model.bin")

    def classify(self, text: str, segment=False):
        '''
        对输入新闻进行分类

        返回值形如：tuple(('__label__运动',), array([0.97278696]))
        '''

        if segment == False:
            if self.wa is None:
                self.wa = WordAnalyzer()
            word_list = self.wa.seg(text)
            return self.model.predict(" ".join(word_list))
        else:
            return self.model.predict(text)
